package org.lanyu.springainovel.common.config;

import lombok.extern.slf4j.Slf4j;
import org.lanyu.springainovel.common.entity.AiModel;
import org.springframework.ai.embedding.EmbeddingModel;
import org.springframework.ai.ollama.OllamaEmbeddingModel;
import org.springframework.ai.ollama.api.OllamaApi;
import org.springframework.ai.ollama.api.OllamaOptions;
import org.springframework.ai.openai.OpenAiEmbeddingModel;
import org.springframework.ai.openai.OpenAiEmbeddingOptions;
import org.springframework.ai.openai.api.OpenAiApi;
import org.springframework.stereotype.Component;

import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

/**
 * 嵌入模型管理�?
 * 用于动态创建和管理不同类型的嵌入模�?
 */
@Slf4j
@Component
public class EmbeddingModelManager {
    
    private final Map<Long, EmbeddingModel> embeddingModelCache = new ConcurrentHashMap<>();
    
    /**
     * 根据模型配置创建嵌入模型
     * 
     * @param modelConfig 模型配置
     * @return 嵌入模型实例
     */
    public EmbeddingModel createEmbeddingModel(AiModel modelConfig) {
        if (modelConfig == null) {
            throw new IllegalArgumentException("模型配置不能为空");
        }
        
        try {
            EmbeddingModel embeddingModel = switch (modelConfig.getVendor().toLowerCase()) {
                case "ollama" -> createOllamaEmbeddingModel(modelConfig);
                default ->  createOpenAiEmbeddingModel(modelConfig);
            };
            
            log.info("创建嵌入模型成功: {} - {}", modelConfig.getName(), modelConfig.getType());
            return embeddingModel;
        } catch (Exception e) {
            log.error("创建嵌入模型失败: {} - {}", modelConfig.getName(), e.getMessage());
            throw new RuntimeException("创建嵌入模型失败", e);
        }
    }
    
    /**
     * 获取或创建嵌入模型（带缓存）
     * 
     * @param modelConfig 模型配置
     * @return 嵌入模型实例
     */
    public EmbeddingModel getOrCreateEmbeddingModel(AiModel modelConfig) {
        return embeddingModelCache.computeIfAbsent(modelConfig.getId(), 
            k -> createEmbeddingModel(modelConfig));
    }
    
    /**
     * 清除指定模型的缓�?
     * 
     * @param modelId 模型ID
     */
    public void clearCache(Long modelId) {
        embeddingModelCache.remove(modelId);
        log.info("清除嵌入模型缓存: {}", modelId);
    }
    
    /**
     * 清除所有缓�?
     */
    public void clearAllCache() {
        embeddingModelCache.clear();
        log.info("清除所有嵌入模型缓存");
    }
    
    /**
     * 创建Ollama嵌入模型
     */
    private EmbeddingModel createOllamaEmbeddingModel(AiModel modelConfig) {
        OllamaApi ollamaApi = OllamaApi.builder()
            .baseUrl(modelConfig.getApiUrl())
            .build();
        
        return OllamaEmbeddingModel.builder()
            .ollamaApi(ollamaApi)
            .defaultOptions(OllamaOptions.builder()
                .model(modelConfig.getName())
                .build())
            .build();
    }
    
    /**
     * 创建OpenAI嵌入模型
     */
    private EmbeddingModel createOpenAiEmbeddingModel(AiModel modelConfig) {
        OpenAiApi openAiApi = OpenAiApi.builder()
            .baseUrl(modelConfig.getApiUrl())
            .apiKey(modelConfig.getApiKey())
            .build();
        
        return new OpenAiEmbeddingModel(openAiApi, 
            org.springframework.ai.document.MetadataMode.EMBED,
            OpenAiEmbeddingOptions.builder()
                .model(modelConfig.getName())
                .build());
    }
} 
